#!/usr/bin/env python3
import sys
import argparse
import onnx
from pathlib import Path

# 设置包搜索路径
__p = Path(__file__).parent
sys.path.append(str(__p))

from py_lib.onnx_op_wrapper.round_clip_fix_wrapper import RoundClipFixWrapper
from py_lib.onnx_op_wrapper.element_wise_op_wrapper import ElementWiseOpWrappper

OP_RoundClipFix='RoundClipFix'
OP_Abs='Abs'
OP_MyAbs='MyAbs'
SUPPORT_OPS = [OP_RoundClipFix, OP_Abs, OP_MyAbs]

def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '-op',
        '--operator',
        choices=SUPPORT_OPS, 
        default=OP_RoundClipFix,
        type=str,
        help=f'要包装的算子类型,支持的类型{SUPPORT_OPS},缺省{OP_RoundClipFix}'
    )

    parser.add_argument(
        '-s',
        '--save',
        type=str,
        default='tmp/tmp.onnx',
        help='模型的保存路径,缺省为tmp/tmp.onnx'
    )

    return parser.parse_args()

def main():
    args = parse_arguments()

    if args.operator == OP_RoundClipFix:
        wrapper = RoundClipFixWrapper(19)
    elif args.operator == OP_Abs:
        wrapper = ElementWiseOpWrappper(OP_Abs)
    elif args.operator == OP_MyAbs:
        wrapper = ElementWiseOpWrappper(OP_MyAbs)
    else:
        print(f'不支持的算子{args.operator}')
        sys.exit(1)

    model = wrapper.build_model()
    # 算子未实行，不能check
    # onnx.checker.check_model(model)
    onnx.save(model, args.save)

if __name__ == '__main__':
    main()
